import tkinter as tk
from tkinter import ttk
import random
import math
from ui.BasePage import BasePage

class ModelTrainingPage(BasePage):
    """模型训练页面（整合版）"""
    def __init__(self, parent, controller):
        super().__init__(parent, controller)
        
        # 左侧训练状态和按钮
        self.left_frame = ttk.Frame(self.frame)
        self.left_frame.pack(side=tk.LEFT, fill=tk.BOTH, padx=10, pady=5, expand=True)

        # 训练状态文本框
        train_status = ttk.LabelFrame(self.left_frame, text="训练状态")
        train_status.pack(fill=tk.X, padx=5, pady=5)
        self.status_text = tk.Text(train_status, height=10, wrap=tk.WORD)
        self.status_text.pack(fill=tk.X, padx=5, pady=5)
        self.status_text.insert(tk.END, "等待开始训练...")
        self.status_text.config(state=tk.DISABLED)

        # 操作按钮
        self.btn_frame = ttk.Frame(self.left_frame)
        self.btn_frame.pack(fill=tk.X, padx=5, pady=5)
        self.start_btn = ttk.Button(self.btn_frame, text="开始训练", style="Info.TButton",
                                   command=self.start_training)
        self.start_btn.pack(fill=tk.X, padx=5, pady=2)
        self.stop_btn = ttk.Button(self.btn_frame, text="停止训练", style="Info.TButton", 
                                  state=tk.DISABLED, command=self.stop_training)
        self.stop_btn.pack(fill=tk.X, padx=5, pady=2)
        self.save_btn = ttk.Button(self.btn_frame, text="保存模型", style="Info.TButton",
                                  command=self.save_model)
        self.save_btn.pack(fill=tk.X, padx=5, pady=2)
        self.load_btn = ttk.Button(self.btn_frame, text="加载模型", style="Info.TButton",
                                  command=self.load_model)
        self.load_btn.pack(fill=tk.X, padx=5, pady=2)

        # 右侧图表部分
        self.right_frame = ttk.Frame(self.frame)
        self.right_frame.pack(side=tk.RIGHT, fill=tk.BOTH, padx=10, pady=5, expand=True)
        self.canvas = tk.Canvas(self.right_frame, bg="white")
        self.canvas.pack(fill=tk.BOTH, expand=True)

        # 绑定窗口大小变化事件
        self.frame.bind("<Configure>", self.on_resize)
        
        # 初始化训练状态
        self.is_training = False
        self.epochs = 50
        self.current_epoch = 0
        self.train_losses = []
        self.val_losses = []
        
        # 初始绘制图表
        self.draw_loss_chart()

    def on_resize(self, event):
        """窗口大小变化时重绘图表"""
        self.draw_loss_chart()

    def draw_loss_chart(self):
        """绘制损失曲线图表"""
        self.canvas.delete("all")
        canvas_width = self.canvas.winfo_width()
        canvas_height = self.canvas.winfo_height()
        padding = max(20, canvas_width // 20)  # 动态边距
        
        # 绘制图表标题
        self.canvas.create_text(canvas_width // 2, padding // 2, 
                               text="Training Visualization: Loss Curves", 
                               font=("微软雅黑", 10, "bold"))
        
        # 绘制图例
        legend_y = padding + 5
        self.canvas.create_rectangle(canvas_width - 150, legend_y, canvas_width - 130, legend_y + 20, fill="blue")
        self.canvas.create_text(canvas_width - 120, legend_y + 10, text="训练损失")
        self.canvas.create_rectangle(canvas_width - 150, legend_y + 25, canvas_width - 130, legend_y + 45, fill="red")
        self.canvas.create_text(canvas_width - 120, legend_y + 35, text="验证损失")
        
        # 生成模拟数据（仅首次绘制时）
        if not self.train_losses:
            self.train_losses = [0.8 * math.exp(-i/10) + random.uniform(0, 0.1) for i in range(self.epochs)]
            self.val_losses = [0.6 * math.exp(-i/10) + random.uniform(0, 0.05) for i in range(self.epochs)]
        
        # 计算缩放比例
        max_loss = max(max(self.train_losses), max(self.val_losses)) if self.train_losses else 1
        x_scale = (canvas_width - 2 * padding) / self.epochs
        y_scale = (canvas_height - 2 * padding) / max_loss if max_loss > 0 else 1
        
        # 绘制坐标轴
        self.canvas.create_line(padding, canvas_height - padding, 
                               canvas_width - padding, canvas_height - padding, width=2)  # x轴
        self.canvas.create_line(padding, padding, padding, canvas_height - padding, width=2)  # y轴
        self.canvas.create_text(canvas_width - padding, canvas_height - padding + 10, 
                               text="迭代次数", font=("微软雅黑", 8))
        self.canvas.create_text(padding - 10, canvas_height - padding, 
                               text="损失", font=("微软雅黑", 8), angle=90)
        
        # 绘制训练损失曲线
        points = []
        for i, loss in enumerate(self.train_losses):
            x = padding + i * x_scale
            y = canvas_height - padding - loss * y_scale
            points.append(x)
            points.append(y)
        if points:
            self.canvas.create_line(points, fill="blue", width=2)
        
        # 绘制验证损失曲线
        points = []
        for i, loss in enumerate(self.val_losses):
            x = padding + i * x_scale
            y = canvas_height - padding - loss * y_scale
            points.append(x)
            points.append(y)
        if points:
            self.canvas.create_line(points, fill="red", width=2)

    def update_training_status(self, epoch, train_loss, val_loss, is_complete=False):
        """更新训练状态显示"""
        status = f"Epoch {epoch+1}/{self.epochs} - Train Loss: {train_loss:.4f}, Train Acc: {1 - train_loss:.4f}, "
        status += f"Val Loss: {val_loss:.4f}, Val Acc: {1 - val_loss:.4f}"
        
        self.status_text.config(state=tk.NORMAL)
        if is_complete:
            self.status_text.delete(1.0, tk.END)
            self.status_text.insert(tk.END, status + "\n训练完成！\n")
            self.status_text.insert(tk.END, f"最佳验证损失: {min(self.val_losses):.4f}")
        else:
            if self.status_text.get(1.0, tk.END).strip() == "等待开始训练...":
                self.status_text.delete(1.0, tk.END)
            self.status_text.insert(tk.END, status + "\n")
        
        self.status_text.see(tk.END)  # 滚动到底部
        self.status_text.config(state=tk.DISABLED)

    def start_training(self):
        """开始训练过程"""
        self.is_training = True
        self.start_btn.config(state=tk.DISABLED)
        self.stop_btn.config(state=tk.NORMAL)
        
        # 模拟训练过程（实际应用中应替换为真实训练逻辑）
        self.current_epoch = 0
        self.train_losses = []
        self.val_losses = []
        
        self.status_text.config(state=tk.NORMAL)
        self.status_text.delete(1.0, tk.END)
        self.status_text.config(state=tk.DISABLED)
        
        self.perform_training_step()

    def perform_training_step(self):
        """执行单步训练（模拟）"""
        if not self.is_training or self.current_epoch >= self.epochs:
            self.is_training = False
            self.start_btn.config(state=tk.NORMAL)
            self.stop_btn.config(state=tk.DISABLED)
            return
        
        # 模拟训练损失（指数衰减+随机波动）
        train_loss = 0.8 * math.exp(-self.current_epoch/10) + random.uniform(0, 0.1)
        val_loss = 0.6 * math.exp(-self.current_epoch/10) + random.uniform(0, 0.05)
        
        self.train_losses.append(train_loss)
        self.val_losses.append(val_loss)
        self.update_training_status(self.current_epoch, train_loss, val_loss)
        self.draw_loss_chart()  # 更新图表
        
        self.current_epoch += 1
        
        if self.current_epoch < self.epochs:
            # 模拟训练间隔（实际应用中应使用线程处理）
            self.frame.after(100, self.perform_training_step)
        else:
            self.update_training_status(self.current_epoch-1, train_loss, val_loss, is_complete=True)
            self.is_training = False
            self.start_btn.config(state=tk.NORMAL)
            self.stop_btn.config(state=tk.DISABLED)

    def stop_training(self):
        """停止训练"""
        self.is_training = False
        self.start_btn.config(state=tk.NORMAL)
        self.stop_btn.config(state=tk.DISABLED)
        self.status_text.config(state=tk.NORMAL)
        self.status_text.insert(tk.END, "\n训练已停止\n")
        self.status_text.config(state=tk.DISABLED)

    def save_model(self):
        """保存模型（示例功能）"""
        self.status_text.config(state=tk.NORMAL)
        self.status_text.insert(tk.END, "\n模型已保存至 model_weights.pth\n")
        self.status_text.config(state=tk.DISABLED)

    def load_model(self):
        """加载模型（示例功能）"""
        self.status_text.config(state=tk.NORMAL)
        self.status_text.insert(tk.END, "\n已加载预训练模型 model_weights.pth\n")
        self.status_text.config(state=tk.DISABLED)